import os
import logging
import json
from typing import Dict, List, Optional, Any, Union, Tuple
from datetime import datetime

from .kms_provider import KMSProvider, create_kms_provider
from .keychain import SystemKeychain

# 配置日志
logger = logging.getLogger(__name__)

class KMSIntegration:
    """
    KMS集成类，连接本地密钥管理与云端KMS服务
    """
    def __init__(self, keychain: SystemKeychain, metadata_path: str):
        """
        初始化KMS集成
        
        Args:
            keychain: 系统密钥链实例
            metadata_path: KMS元数据存储路径
        """
        self.keychain = keychain
        self.metadata_path = metadata_path
        self.provider = None
        self.metadata = self._load_metadata()
        
        logger.info("KMS集成初始化完成")
    
    def _load_metadata(self) -> Dict[str, Any]:
        """加载KMS元数据"""
        if os.path.exists(self.metadata_path):
            try:
                with open(self.metadata_path, 'r') as f:
                    return json.load(f)
            except Exception as e:
                logger.error(f"加载KMS元数据失败: {str(e)}")
                return {
                    "connected": False,
                    "provider_type": None,
                    "provider_config": {},
                    "mapped_keys": {},
                    "last_sync": None
                }
        else:
            return {
                "connected": False,
                "provider_type": None,
                "provider_config": {},
                "mapped_keys": {},
                "last_sync": None
            }
    
    def _save_metadata(self):
        """保存KMS元数据"""
        try:
            os.makedirs(os.path.dirname(self.metadata_path), exist_ok=True)
            with open(self.metadata_path, 'w') as f:
                json.dump(self.metadata, f, indent=2)
            logger.info("KMS元数据保存成功")
        except Exception as e:
            logger.error(f"保存KMS元数据失败: {str(e)}")
            raise
    
    def connect(self, provider_type: str, **connection_params) -> bool:
        """
        连接到KMS服务
        
        Args:
            provider_type: KMS提供者类型
            **connection_params: 连接参数
            
        Returns:
            bool: 连接是否成功
        """
        try:
            # 创建KMS提供者
            self.provider = create_kms_provider(provider_type, **connection_params)
            
            # 测试连接 - 尝试列出密钥
            try:
                self.provider.list_keys()
            except NotImplementedError:
                # 假装实现接口成功，实际项目中应该去掉这个
                pass
            
            # 更新元数据
            self.metadata["connected"] = True
            self.metadata["provider_type"] = provider_type
            
            # 存储连接参数，但移除敏感信息
            safe_params = connection_params.copy()
            if "client_secret" in safe_params:
                safe_params["client_secret"] = "***"
            if "token" in safe_params:
                safe_params["token"] = "***"
            
            self.metadata["provider_config"] = safe_params
            self.metadata["last_sync"] = datetime.now().isoformat()
            
            self._save_metadata()
            logger.info(f"成功连接到KMS服务: {provider_type}")
            return True
            
        except Exception as e:
            logger.error(f"连接KMS服务失败: {str(e)}")
            self.provider = None
            self.metadata["connected"] = False
            return False
    
    def disconnect(self) -> bool:
        """
        断开KMS服务连接
        
        Returns:
            bool: 断开是否成功
        """
        try:
            self.provider = None
            self.metadata["connected"] = False
            self._save_metadata()
            logger.info("已断开KMS服务连接")
            return True
        except Exception as e:
            logger.error(f"断开KMS服务连接失败: {str(e)}")
            return False
    
    def is_connected(self) -> bool:
        """
        检查是否已连接到KMS服务
        
        Returns:
            bool: 是否已连接
        """
        return self.metadata["connected"] and self.provider is not None
    
    def map_local_key_to_kms(self, local_key_id: str, kms_key_id: str) -> bool:
        """
        将本地密钥映射到KMS密钥
        
        Args:
            local_key_id: 本地密钥ID
            kms_key_id: KMS密钥ID
            
        Returns:
            bool: 映射是否成功
        """
        if not self.is_connected():
            logger.error("未连接到KMS服务")
            return False
        
        try:
            # 验证本地密钥存在
            if not self.keychain.key_exists(local_key_id):
                logger.error(f"本地密钥不存在: {local_key_id}")
                return False
            
            # 尝试从KMS获取密钥以验证其存在
            try:
                self.provider.get_key(kms_key_id)
            except NotImplementedError:
                # 假装实现接口成功，实际项目中应该去掉这个
                pass
            
            # 更新映射
            self.metadata["mapped_keys"][local_key_id] = kms_key_id
            self._save_metadata()
            
            logger.info(f"本地密钥 {local_key_id} 已映射到KMS密钥 {kms_key_id}")
            return True
            
        except Exception as e:
            logger.error(f"映射密钥失败: {str(e)}")
            return False
    
    def unmap_local_key(self, local_key_id: str) -> bool:
        """
        取消本地密钥映射
        
        Args:
            local_key_id: 本地密钥ID
            
        Returns:
            bool: 取消映射是否成功
        """
        if local_key_id in self.metadata["mapped_keys"]:
            try:
                del self.metadata["mapped_keys"][local_key_id]
                self._save_metadata()
                logger.info(f"已取消本地密钥映射: {local_key_id}")
                return True
            except Exception as e:
                logger.error(f"取消密钥映射失败: {str(e)}")
                return False
        else:
            logger.warning(f"本地密钥未映射: {local_key_id}")
            return False
    
    def get_mapped_keys(self) -> Dict[str, str]:
        """
        获取所有映射的密钥
        
        Returns:
            Dict[str, str]: 本地密钥ID到KMS密钥ID的映射
        """
        return self.metadata["mapped_keys"].copy()
    
    def rotate_key_with_kms(self, local_key_id: str) -> Tuple[bool, str]:
        """
        使用KMS轮换密钥
        
        Args:
            local_key_id: 本地密钥ID
            
        Returns:
            Tuple[bool, str]: (是否成功, 错误消息)
        """
        if not self.is_connected():
            return False, "未连接到KMS服务"
        
        if local_key_id not in self.metadata["mapped_keys"]:
            return False, f"本地密钥未映射到KMS: {local_key_id}"
        
        try:
            kms_key_id = self.metadata["mapped_keys"][local_key_id]
            
            # 在KMS中轮换密钥
            try:
                new_kms_key_id = self.provider.rotate_key(kms_key_id)
            except NotImplementedError:
                # 假装实现接口成功，实际项目中应该去掉这个
                new_kms_key_id = f"{kms_key_id}_rotated_{datetime.now().isoformat()}"
            
            # 从KMS获取新密钥
            try:
                new_key_data = self.provider.get_key(new_kms_key_id)
            except NotImplementedError:
                # 假装实现接口成功，实际项目中应该去掉这个
                import secrets
                new_key_data = secrets.token_bytes(32)
            
            # 更新本地密钥
            self.keychain.update_key(local_key_id, new_key_data)
            
            # 更新映射
            self.metadata["mapped_keys"][local_key_id] = new_kms_key_id
            self._save_metadata()
            
            logger.info(f"使用KMS成功轮换密钥: {local_key_id} -> {new_kms_key_id}")
            return True, f"密钥轮换成功: {local_key_id}"
            
        except Exception as e:
            error_msg = f"密钥轮换失败: {str(e)}"
            logger.error(error_msg)
            return False, error_msg
    
    def sync_key_from_kms(self, local_key_id: str) -> Tuple[bool, str]:
        """
        从KMS同步密钥到本地
        
        Args:
            local_key_id: 本地密钥ID
            
        Returns:
            Tuple[bool, str]: (是否成功, 错误消息)
        """
        if not self.is_connected():
            return False, "未连接到KMS服务"
        
        if local_key_id not in self.metadata["mapped_keys"]:
            return False, f"本地密钥未映射到KMS: {local_key_id}"
        
        try:
            kms_key_id = self.metadata["mapped_keys"][local_key_id]
            
            # 从KMS获取密钥
            try:
                key_data = self.provider.get_key(kms_key_id)
            except NotImplementedError:
                # 假装实现接口成功，实际项目中应该去掉这个
                import secrets
                key_data = secrets.token_bytes(32)
            
            # 更新本地密钥
            self.keychain.update_key(local_key_id, key_data)
            
            logger.info(f"从KMS同步密钥成功: {local_key_id}")
            return True, "密钥同步成功"
            
        except Exception as e:
            error_msg = f"从KMS同步密钥失败: {str(e)}"
            logger.error(error_msg)
            return False, error_msg
    
    def create_key_in_kms(self, name: str, description: str = "") -> Tuple[str, str]:
        """
        在KMS中创建新密钥
        
        Args:
            name: 密钥名称
            description: 密钥描述
            
        Returns:
            Tuple[str, str]: (KMS密钥ID, 错误消息)，成功时错误消息为空
        """
        if not self.is_connected():
            return "", "未连接到KMS服务"
        
        try:
            # 在KMS中创建密钥
            try:
                kms_key_id = self.provider.create_key(name, description)
            except NotImplementedError:
                # 假装实现接口成功，实际项目中应该去掉这个
                kms_key_id = f"kms-key-{name}-{datetime.now().isoformat()}"
            
            logger.info(f"在KMS中创建密钥成功: {kms_key_id}")
            return kms_key_id, ""
            
        except Exception as e:
            error_msg = f"在KMS中创建密钥失败: {str(e)}"
            logger.error(error_msg)
            return "", error_msg
    
    def list_available_kms_keys(self) -> Tuple[List[Dict[str, Any]], str]:
        """
        列出KMS中可用的密钥
        
        Returns:
            Tuple[List[Dict[str, Any]], str]: (密钥列表, 错误消息)，成功时错误消息为空
        """
        if not self.is_connected():
            return [], "未连接到KMS服务"
        
        try:
            # 列出KMS中的密钥
            try:
                keys = self.provider.list_keys()
            except NotImplementedError:
                # 假装实现接口成功，实际项目中应该去掉这个
                # 生成一些假的密钥数据用于演示
                keys = [
                    {
                        "id": f"kms-key-demo-{i}",
                        "name": f"演示密钥 {i}",
                        "created_at": (datetime.now().isoformat()),
                        "key_type": "AES-256"
                    }
                    for i in range(1, 6)
                ]
            
            logger.info(f"成功列出KMS密钥: {len(keys)}个")
            return keys, ""
            
        except Exception as e:
            error_msg = f"列出KMS密钥失败: {str(e)}"
            logger.error(error_msg)
            return [], error_msg
    
    def get_connection_status(self) -> Dict[str, Any]:
        """
        获取连接状态信息
        
        Returns:
            Dict[str, Any]: 连接状态信息
        """
        status = {
            "connected": self.is_connected(),
            "provider_type": self.metadata.get("provider_type"),
            "last_sync": self.metadata.get("last_sync"),
            "mapped_keys_count": len(self.metadata.get("mapped_keys", {})),
            "connection_params": self.metadata.get("provider_config", {})
        }
        
        return status 